package com.example.JavaWebApi.core.apifilters.components;

import lombok.extern.slf4j.Slf4j;
import org.apache.catalina.connector.CoyoteOutputStream;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintWriter;

/**
 * @program: cnu-api-java
 * @description:
 * @author: 杨泽
 * @create: 2023-07-26 13:52
 */
@Slf4j
public class ResponseWrapper extends HttpServletResponseWrapper {
    private final ByteArrayOutputStream buffer;
    private final ServletOutputStream out;
    private final PrintWriter writer;


    public ResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);

        buffer = new ByteArrayOutputStream(2048);
        out = new WrapperOutputStream(buffer, response);
        writer = new WrapperWriter(buffer, response);
    }

    //重载父类获取outputstream的方法
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return out;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return writer;
    }


    @Override
    public void reset() {
        buffer.reset();
    }

    public byte[] getResponseData() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }

    public String getContent() throws IOException {
        flushBuffer();
        return buffer.toString();
    }


    //内部类，对ServletOutputStream进行包装，指定输出流的输出端

    public class WrapperOutputStream extends ServletOutputStream {
        private OutputStream innerOut;
        private HttpServletResponse response;

        public WrapperOutputStream(OutputStream innerOut, HttpServletResponse response) {
            super();
            this.response = response;
            this.innerOut = innerOut;
        }

        @Override
        public boolean isReady() {
            if (response == null) {
                return false;
            }
            try {
                return response.getOutputStream().isReady();
            } catch (IOException e) {
                e.printStackTrace();
            }
            return false;
        }

        @Override
        public void setWriteListener(WriteListener listener) {
            if (response != null) {
                try {
                    ((CoyoteOutputStream) response.getOutputStream()).setWriteListener(listener);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

        }

        //关键在这里
        @Override
        public void write(int b) throws IOException {
            if (response != null) {
                response.getOutputStream().write(b);
            }
            innerOut.write(b);
        }
    }

    public class WrapperWriter extends PrintWriter {
        private HttpServletResponse response;
        ByteArrayOutputStream output;

        public WrapperWriter(ByteArrayOutputStream out, HttpServletResponse response) {
            super(out);
            this.response = response;
            this.output = out;
        }

        //关键在这里
        @Override
        public void write(int b) {
            super.write(b);
            try {
                response.getWriter().write(b);
            } catch (IOException e) {
                e.printStackTrace();
                this.setError();
            }
        }

        //关键在这里
        @Override
        public void write(String s, int off, int len) {
            super.write(s, off, len);
            try {
                response.getWriter().write(s, off, len);
            } catch (IOException e) {
                e.printStackTrace();
                this.setError();
            }
        }
    }
}